import socket
import json
from app import DocumentVectorizer, SearchService
from utils.logger import logger
from config import Config
import threading
import requests
import os

class TCPServer:
    def __init__(self):
        self.host = Config.TCP_SERVER_HOST
        self.port = Config.TCP_PORT
        self.buffer_size = Config.TCP_BUFFER_SIZE
        self.server_socket = None
        self.running = False
        self.thread = None
        self.vectorizer = DocumentVectorizer()
        self.search_service = SearchService()
        
        # 启动时加载索引
        self.load_index()
        
    def load_index(self):
        """加载索引"""
        try:
            if os.path.exists(Config.DB_DIR):
                logger.info(f"正在加载索引: {Config.DB_DIR}")
                self.search_service.load_index(Config.DB_DIR)
                logger.info("索引加载成功")
            else:
                logger.warning("未找到索引目录，将使用空索引启动")
                
        except Exception as e:
            logger.error(f"加载索引失败: {str(e)}")
            logger.warning("将使用空索引启动")
        
    def start(self):
        """启动TCP服务器"""
        try:
            # 创建服务器套接字
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            
            # 绑定地址前先关闭可能存在的连接
            try:
                self.server_socket.shutdown(socket.SHUT_RDWR)
            except:
                pass
            
            try:
                self.server_socket.close()
            except:
                pass
                
            # 重新绑定和监听
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            
            logger.info(f"正在绑定地址 {self.host}:{self.port}")
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(5)
            
            logger.info(f"TCP服务器启动在 {self.host}:{self.port}")
            
            # 设置运行标志
            self.running = True
            
            # 在新线程中运行服务
            self.thread = threading.Thread(target=self._run_server)
            self.thread.daemon = True
            self.thread.start()
            
        except Exception as e:
            logger.error(f"启动TCP服务器失败: {str(e)}")
            raise
    
    def stop(self):
        """停止TCP服务器"""
        try:
            self.running = False
            
            # 关闭服务器套接字
            if self.server_socket:
                self.server_socket.close()
                self.server_socket = None
            
            # 等待服务线程结束
            if self.thread and self.thread.is_alive():
                self.thread.join(timeout=5)
                
            logger.info("TCP服务器已停止")
            
        except Exception as e:
            logger.error(f"停止TCP服务器时发生错误: {str(e)}")
    
    def _run_server(self):
        """服务器主循环"""
        while self.running:
            try:
                # 设置超时以便能够检查running标志
                self.server_socket.settimeout(1)
                
                try:
                    client_socket, address = self.server_socket.accept()
                    logger.info(f"接受来自 {address} 的连接")
                    
                    # 在新线程中处理客户端连接
                    client_thread = threading.Thread(
                        target=self._handle_client,
                        args=(client_socket, address)
                    )
                    client_thread.daemon = True
                    client_thread.start()
                    
                except socket.timeout:
                    continue
                    
            except Exception as e:
                if self.running:  # 只在服务器运行时记录错误
                    logger.error(f"接受连接失败: {str(e)}")
    
    def _handle_client(self, client_socket, address):
        """处理客户端连接"""
        try:
            while self.running:
                # 接收数据
                data = client_socket.recv(self.buffer_size)
                if not data:
                    break
                
                # 处理数据
                try:
                    # 这里添加您的数据处理逻辑
                    response = self._process_data(data)
                    
                    # 发送响应
                    client_socket.send(response.encode())
                    
                except Exception as e:
                    logger.error(f"处理客户端数据失败: {str(e)}")
                    error_response = json.dumps({"error": str(e)})
                    client_socket.send(error_response.encode())
                
        except Exception as e:
            logger.error(f"处理客户端 {address} 时发生错误: {str(e)}")
        
        finally:
            client_socket.close()
            logger.info(f"客户端 {address} 连接已关闭")
    
    def _process_data(self, data):
        """处理接收到的数据"""
        try:
            # 解码并解析JSON数据
            message = json.loads(data.decode())
            
            # 这里添加您的业务逻辑处理代码
            cmd_type = message.get('type', '')
            
            if cmd_type == 'add_documents':
                # 处理添加文档请求
                url = message.get('data', '')  # 获取URL而不是目录路径
                response = self.handle_add_documents(url)
                
            elif cmd_type == 'search':
                # 处理检索请求
                query = message.get('data', '')
                top_k = message.get('top_k', Config.DEFAULT_TOP_K)
                response = self.handle_search(query, top_k)
                
            elif cmd_type == 'get_documents':
                # 获取文档内容
                indices = message.get('indices', [])
                response = self.handle_get_documents(indices)
                
            else:
                response = {
                    'status': 0,
                    'type': "未知错误",
                    'data': "未知错误"
                }
            
            # 返回处理结果（添加换行符）
            return json.dumps(response) + '\n'
            
        except json.JSONDecodeError:
            return json.dumps({"error": "Invalid JSON format"}) + '\n'
        except Exception as e:
            return json.dumps({"error": str(e)}) + '\n'
    
    def handle_add_documents(self, url):
        """处理添加文档请求，从URL下载文档"""
        try:
            # 创建临时目录
            temp_dir = os.path.join(os.getcwd(), "temp")
            # 如果目录存在,先清除
            if os.path.exists(temp_dir):
                import shutil
                shutil.rmtree(temp_dir)
            os.makedirs(temp_dir)
            
            try:
                # 下载文件
                logger.info(f"开始从 {url} 下载文档")
                response = requests.get(url, stream=True)
                response.raise_for_status()
                
                # 获取文件名
                filename = url.split('/')[-1]
                if not filename:
                    filename = 'downloaded_file.zip'
                    
                # 保存文件
                file_path = os.path.join(temp_dir, filename)
                with open(file_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                            
                logger.info(f"文档下载完成: {file_path}")
                
                # 如果是压缩文件，解压
                if filename.endswith('.zip'):
                    import zipfile
                    with zipfile.ZipFile(file_path, 'r') as zip_ref:
                        zip_ref.extractall(temp_dir)
                    logger.info("解压完成")
                    
                # 处理文档
                chunks, qa_pairs, embeddings, doc_types = self.vectorizer.vectorize_documents(temp_dir)
                if not chunks or not embeddings:
                    return {'error': '文档处理失败'}
                    
                # 添加到检索系统
                self.search_service.add_documents(chunks, qa_pairs, embeddings, doc_types)
                
                # 自动保存索引
                os.makedirs(Config.DB_DIR, exist_ok=True)
                self.search_service.save_index(Config.DB_DIR)
                logger.info(f"索引已保存到: {Config.DB_DIR}")
                
                return {
                    'status': 1,
                    'type': "add_documents",
                    'data': "知识库接收成功"
                }
                
            finally:
                pass
                # 清理临时文件
                # import shutil
                # shutil.rmtree(temp_dir)
                
        except Exception as e:
            logger.error(f"处理文档失败: {str(e)}")
            return {
                    'status': 0,
                    'type': "add_documents",
                    'data': "知识库下载失败"
                }
            
    def handle_search(self, query, top_k):
        """处理检索请求"""
        try:
            logger.info(f"开始检索，query: {query}, top_k: {top_k}")
            context = ""
            if len(query) > 3:
                results = self.search_service.search(query, top_k)
                
                # 获取文档内容
                indices = [r['chunk_index'] for r in results]
                documents = self.search_service.retriever.get_documents_by_indices(indices)
                
                # 组装结果
                # final_results = []
                # for doc, result in zip(documents, results):
                    # final_results.append({
                    #     'content': doc,
                    #     'score': result['score']
                    # })
                context = '##'.join(documents).replace('\n', ' ')
            if top_k == 1:
                return {
                        'status': 1,
                        'type': "search",
                        'question': query,
                        'data': context,
                        'same_answer': 1
                    }
            else:
                return {
                    'status': 1,
                    'type': "search",
                    'question': query,
                    'data': context,
                    'same_answer': 0
                }
        except Exception as e:
            logger.error(f"检索失败: {str(e)}")
            return {
                    'status': 0,
                    'type': "search",
                    'question': query,
                    'data': "读取失败",
                    'same_answer': 0
                }
            
    def handle_get_documents(self, indices):
        """处理获取文档内容请求"""
        try:
            documents = self.search_service.retriever.get_documents_by_indices(indices)
            return {
                'message': 'success',
                'documents': documents
            }
        except Exception as e:
            return {'error': str(e)}

if __name__ == "__main__":
    server = TCPServer()
    server.start() 